import os
import pandas as pd
import matplotlib.pyplot as plt
from math import *
from cmath import phase
from pandas import DataFrame as df

T=1000*4*pi/1e7

plt.rcParams['font.size']=11
plt.rcParams['font.family']="Inconsolata"

path=os.path.join("datasetsedd","test")
fs=[fs for parent,dirnames,fs in os.walk(path)][0]
fs=[os.path.join(path,det) for det in fs if ".dat" in det]

for _,each in enumerate(fs[:1]):
    print(each)
    basename=os.path.splitext(os.path.basename(each))[0]
    content=[]
    with open(os.path.join(path,basename,"forward","temp.res"),"r") as fp:
        content.extend([det.split() for det in fp])

    swap=content[:]

    content=swap[:int(len(swap)/2)] 
    header=content[:8]
    data=df(content[8:],columns=header[1][1:])
    data["Real"]=data["Real"].astype(float)
    data["Imag"]=data["Imag"].astype(float)

    TEFWD=df()
    for index,row in data.iloc[:].iterrows():
        zxyr,zxyi=row["Real"]/T,row["Imag"]/T
        zxy=zxyr*zxyr+zxyi*zxyi
        rxy=0.2*float(row["Period(s)"])*zxy
        pxy=abs(phase(complex(zxyr,zxyi)))*180.0/pi

        temp=df([[row["Code"],row["Period(s)"],rxy,pxy]],columns=["STATION","Period","RXY","PXY"])
        TEFWD=pd.concat([TEFWD,temp],axis=0)
    
    content=swap[int(len(swap)/2):]
    header=content[:8]
    data=df(content[8:],columns=header[1][1:])
    data["Real"]=data["Real"].astype(float)
    data["Imag"]=data["Imag"].astype(float)

    TMFWD=df()
    for index,row in data.iloc[:].iterrows():
        zxyr,zxyi=row["Real"]/T,row["Imag"]/T
        zxy=zxyr*zxyr+zxyi*zxyi
        rxy=0.2*float(row["Period(s)"])*zxy
        pxy=abs(phase(complex(zxyr,zxyi)))*180.0/pi

        temp=df([[row["Code"],row["Period(s)"],rxy,pxy]],columns=["STATION","Period","RXY","PXY"])
        TMFWD=pd.concat([TMFWD,temp],axis=0)

    filenames=[filenames for parent,dirname,filenames in os.walk(os.path.join(path,basename,"inversion"))][0]
    filenames=[det for det in filenames if ".dat" in det and "NLCG" in det]
    filenames.sort(key=lambda x: os.path.getmtime(os.path.join(path,basename,"inversion",x)))
    
    content=[]
    with open(os.path.join(path,basename,"inversion",filenames[-1]),"r") as fp:
        content.extend([det.split() for det in fp])

    swap=content[:]
    content=swap[:int(len(swap)/2)] 
    header=content[:8]
    data=df(content[8:],columns=header[1][1:])
    data["Real"]=data["Real"].astype(float)
    data["Imag"]=data["Imag"].astype(float)

    TEINV=df()
    for index,row in data.iloc[:].iterrows():
        zxyr,zxyi=row["Real"]/T,row["Imag"]/T
        zxy=zxyr*zxyr+zxyi*zxyi
        rxy=0.2*float(row["Period(s)"])*zxy
        pxy=abs(phase(complex(zxyr,zxyi)))*180.0/pi

        temp=df([[row["Code"],row["Period(s)"],rxy,pxy]],columns=["STATION","Period","RXY","PXY"])
        TEINV=pd.concat([TEINV,temp],axis=0)
    
    content=swap[int(len(swap)/2):]
    header=content[:8]
    data=df(content[8:],columns=header[1][1:])
    data["Real"]=data["Real"].astype(float)
    data["Imag"]=data["Imag"].astype(float)

    TMINV=df()
    for index,row in data.iloc[:].iterrows():
        zxyr,zxyi=row["Real"]/T,row["Imag"]/T
        zxy=zxyr*zxyr+zxyi*zxyi
        rxy=0.2*float(row["Period(s)"])*zxy
        pxy=abs(phase(complex(zxyr,zxyi)))*180.0/pi

        temp=df([[row["Code"],row["Period(s)"],rxy,pxy]],columns=["STATION","Period","RXY","PXY"])
        TMINV=pd.concat([TMINV,temp],axis=0)

    if os.path.exists(os.path.join(path,basename,"result.res")):        
        content=[]
        with open(os.path.join(path,basename,"result.res"),"r") as fp:
            content.extend([det.split() for det in fp])

        swap=content[:]
        content=swap[:int(len(swap)/2)] 
        header=content[:8]
        data=df(content[8:],columns=header[1][1:])
        data["Real"]=data["Real"].astype(float)
        data["Imag"]=data["Imag"].astype(float)

        TEIMG=df()
        for index,row in data.iloc[:].iterrows():
            zxyr,zxyi=row["Real"]/T,row["Imag"]/T
            zxy=zxyr*zxyr+zxyi*zxyi
            rxy=0.2*float(row["Period(s)"])*zxy
            pxy=abs(phase(complex(zxyr,zxyi)))*180.0/pi

            temp=df([[row["Code"],row["Period(s)"],rxy,pxy]],columns=["STATION","Period","RXY","PXY"])
            TEIMG=pd.concat([TEIMG,temp],axis=0)
        
        content=swap[int(len(swap)/2):]
        header=content[:8]
        data=df(content[8:],columns=header[1][1:])
        data["Real"]=data["Real"].astype(float)
        data["Imag"]=data["Imag"].astype(float)

        TMIMG=df()
        for index,row in data.iloc[:].iterrows():
            zxyr,zxyi=row["Real"]/T,row["Imag"]/T
            zxy=zxyr*zxyr+zxyi*zxyi
            rxy=0.2*float(row["Period(s)"])*zxy
            pxy=abs(phase(complex(zxyr,zxyi)))*180.0/pi

            temp=df([[row["Code"],row["Period(s)"],rxy,pxy]],columns=["STATION","Period","RXY","PXY"])
            TMIMG=pd.concat([TMIMG,temp],axis=0)
        
        if os.path.exists(os.path.join(path,basename,"compare")) is False:
            os.mkdir(os.path.join(path,basename,"compare"))

        code=TMFWD["STATION"].unique()
        for row in code[:]:
            print(row)
            forward=TMFWD[TMFWD["STATION"]==row]
            image=TMIMG[TMIMG["STATION"]==row]
            inverion=TMINV[TMINV["STATION"]==row]
            freq=list(forward["Period"])
            curve,=plt.plot(freq,list(forward["RXY"]),"b:")
            curve.set_color("black")
            plt.plot(freq,list(image["RXY"]),"r-")
            plt.plot(freq,list(inverion["RXY"]),"b",dashes=[6,2])
            plt.ylim(1,10000.0)
            plt.xscale('log')
            plt.yscale('log')
            plt.xlabel("Period(s)")
            plt.ylabel("res(ohm-m)")
            plt.grid(which="both")
            plt.grid(which='minor',alpha=0.2)
            plt.grid(which='major',alpha=0.5)
            plt.savefig(os.path.join(path,basename,"compare",row+".jpg"),dpi=600)
            plt.cla()
            plt.clf()